"""
简单的增量学习框架使用示例

这个示例展示了如何使用框架的核心组件来实现多任务增量学习。
"""

import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader, TensorDataset

# 导入框架组件
from continuallearning.module import ContinualModule
from continuallearning.callbacks import TaskBoundaryCallback, TaskTrainingCallback
from continuallearning.methods import EWC  # 假设有EWC实现
from continuallearning.config import OptimizerConfig


def create_simple_example():
    """创建一个简单的使用示例"""

    # 1. 创建一个简单的模型
    class SimpleModel(nn.Module):
        def __init__(self, input_dim=10, hidden_dim=64, output_dim=2):
            super().__init__()
            self.fc1 = nn.Linear(input_dim, hidden_dim)
            self.fc2 = nn.Linear(hidden_dim, output_dim)
            self.relu = nn.ReLU()

        def forward(self, x, task_id=None):
            # task_id参数是为了兼容框架接口，这里简单忽略
            x = self.relu(self.fc1(x))
            return self.fc2(x)

    # 2. 创建模型实例
    model = SimpleModel()

    # 3. 配置优化器
    optimizer_cfg = OptimizerConfig(type="Adam", lr=0.001, weight_decay=1e-4)

    # 4. 创建增量学习策略（如果有的话）
    # cl_strategy = EWC(regularization_weight=5000)
    cl_strategy = None  # 暂时不使用策略

    # 5. 创建ContinualModule
    continual_module = ContinualModule(
        model=model, optimizer_cfg=optimizer_cfg, cl_strategy=cl_strategy
    )

    # 6. 创建回调系统
    task_boundary_cb = TaskBoundaryCallback(verbose=True)
    task_training_cb = TaskTrainingCallback(
        task_boundary_callback=task_boundary_cb, reset_optimizer_between_tasks=False
    )

    # 7. 创建Lightning训练器
    trainer = pl.Trainer(
        max_epochs=3,  # 每个任务训练3个epoch
        callbacks=[task_boundary_cb, task_training_cb],
        logger=False,  # 关闭日志以简化输出
        enable_checkpointing=False,
    )

    # 8. 创建模拟的多任务数据
    def create_dummy_data(task_id, batch_size=32, num_samples=100):
        """创建虚拟数据用于演示"""
        # 为每个任务创建稍微不同的数据分布
        X = torch.randn(num_samples, 10) + task_id * 0.5
        y = torch.randint(0, 2, (num_samples,))
        dataset = TensorDataset(X, y)
        return DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # 9. 准备多任务数据
    tasks_data = []
    for task_id in range(3):  # 3个任务
        train_loader = create_dummy_data(task_id, num_samples=100)
        val_loader = create_dummy_data(task_id, num_samples=50)
        tasks_data.append({"train": train_loader, "val": val_loader})

    # 10. 执行多任务训练
    print("开始多任务增量学习训练...")
    results = task_training_cb.train_multiple_tasks(
        trainer=trainer, model=continual_module, tasks_data=tasks_data
    )

    print(f"训练完成！训练了 {len(results)} 个任务")

    # 11. 测试所有任务（可选）
    test_loaders = [create_dummy_data(i, num_samples=30) for i in range(3)]
    test_results = task_training_cb.test_all_tasks(
        trainer=trainer, model=continual_module, test_dataloaders=test_loaders
    )

    print(f"测试完成！测试了 {len(test_results)} 个任务")

    return continual_module, results, test_results


if __name__ == "__main__":
    # 运行示例
    model, train_results, test_results = create_simple_example()
    print("示例运行完成！")
